#pragma once

#include "LocUtils/common/eigen_types.h"
#include "LocUtils/sensor_data/imu_data.hpp"
#include "LocUtils/common/odom.h"
#include "LocUtils/common/pose_data.hpp"
#include "LocUtils/common/math_utils.h"
#include "LocUtils/common/nav_state.h"
#include <glog/logging.h>
#include <iomanip>

namespace LocUtils
{

struct EskfOptions 
{
    EskfOptions() = default;

    /// IMU 测量与零偏参数
    double imu_dt_ = 0.01;  // IMU测量间隔
    // NOTE IMU噪声项都为离散时间，不需要再乘dt，可以由初始化器指定IMU噪声
    double gyro_var_ = 1e-5;       // 陀螺测量标准差
    double acce_var_ = 1e-2;       // 加计测量标准差
    double bias_gyro_var_ = 1e-6;  // 陀螺零偏游走标准差
    double bias_acce_var_ = 1e-4;  // 加计零偏游走标准差

    /// 里程计参数
    double odom_var_ = 0.5;
    double odom_span_ = 0.1;        // 里程计测量间隔
    double wheel_radius_ = 0.155;   // 轮子半径
    double circle_pulse_ = 1024.0;  // 编码器每圈脉冲数

    /// RTK 观测参数
    double gnss_pos_noise_ = 0.1;                   // GNSS位置噪声
    double gnss_height_noise_ = 0.1;                // GNSS高度噪声
    double gnss_ang_noise_ = 1.0 * math::kDEG2RAD;  // GNSS旋转噪声

    /// 激光里程计 信息 观测参数
    double lidar_pos_noise_ = 0.1;                   // 激光位置噪声
    double lidar_height_noise_ = 0.1;                // 激光高度噪声
    double lidar_ang_noise_ = 1.0 * math::kDEG2RAD;  // 激光旋转噪声

    /// 其他配置
    bool update_bias_gyro_ = true;  // 是否更新陀螺bias
    bool update_bias_acce_ = true;  // 是否更新加计bias


};

/**
 * 书本第3章介绍的误差卡尔曼滤波器
 * 可以指定观测GNSS的读数，GNSS应该事先转换到车体坐标系
 *
 * 本书使用18维的ESKF，标量类型可以由S指定，默认取double
 * 变量顺序：p, v, R, bg, ba, grav，与书本对应
 * @tparam S    状态变量的精度，取float或double
 */
template <typename S = double>
class ESKF 
{
    public:
        /// 类型定义
        using SO3 = Sophus::SO3<S>;                     // 旋转变量类型
        using VecT = Eigen::Matrix<S, 3, 1>;            // 向量类型
        using Vec18T = Eigen::Matrix<S, 18, 1>;         // 18维向量类型
        using Mat3T = Eigen::Matrix<S, 3, 3>;           // 3x3矩阵类型
        using MotionNoiseT = Eigen::Matrix<S, 18, 18>;  // 运动噪声类型
        using OdomNoiseT = Eigen::Matrix<S, 3, 3>;      // 里程计噪声类型
        using GnssNoiseT = Eigen::Matrix<S, 6, 6>;      // GNSS噪声类型
        using LidarNoiseT = Eigen::Matrix<S, 6, 6>;      // GNSS噪声类型
        using Mat18T = Eigen::Matrix<S, 18, 18>;        // 18维方差类型

        
        /**
         * 初始零偏取零
         */
        ESKF(EskfOptions option = EskfOptions()) : options_(option) { BuildNoise(option); }

        /**
         * 设置初始条件
         * @param options 噪声项配置
         * @param init_bg 初始零偏 陀螺
         * @param init_ba 初始零偏 加计
         * @param gravity 重力
         */
        void SetInitialConditions(EskfOptions options, const VecT& init_bg, const VecT& init_ba,
                              const VecT& gravity = VecT(0, 0, -9.8)) {
            BuildNoise(options);
            options_ = options;
            bg_ = init_bg;
            ba_ = init_ba;
            g_ = gravity;
            cov_ = Mat18T::Identity() * 1e-4;
        }

        /// 使用IMU递推
        bool Predict(const IMU& imu);

        /// 使用轮速计观测
        // bool ObserveWheelSpeed(const Odom& odom);

        /// 使用Lidar观测
        bool ObserveLidar(const PoseData& pos);

        /**
         * 使用SE3进行观测
         * @param pose  观测位姿
         * @param trans_noise 平移噪声
         * @param ang_noise   角度噪声
         * @return
         */
        bool ObserveSE3(const SE3& pose, double trans_noise = 0.1, double ang_noise = 1.0 * math::kDEG2RAD);
        
        /// 设置状态X
        void SetX(const NavStated& x, const Vec3d& grav = Vec3d(0, 0, -9.8)) 
        {
            current_time_ = x.timestamp_;
            R_ = x.R_;
            p_ = x.p_;
            v_ = x.v_;
            bg_ = x.bg_;
            ba_ = x.ba_;
            g_ = grav;
        }

        /// 设置状态X se3
        void SetX(const PoseData& x, const Vec3d& grav = Vec3d(0, 0, -9.8)) 
        {
            SE3 pos; x.GetSE3(pos);
            R_ = pos.so3();
            p_ = pos.translation();
            g_ = grav;
        }

        void SetX(const SE3& x, const Vec3d& grav = Vec3d(0, 0, -9.8)) 
        {
            R_ = x.so3();
            p_ = x.translation();
            g_ = grav;
        }
        
        /// 获取SE3 状态
        SE3 GetNominalSE3() const { return SE3(R_, p_); }
        
        /// 设置协方差
        void SetCov(const Mat18T& cov) { cov_ = cov; }

        /// 获取重力
        Vec3d GetGravity() const { return g_; }

    private:
        
        void BuildNoise(const EskfOptions& options) {
            double ev = options.acce_var_;
            double et = options.gyro_var_;
            double eg = options.bias_gyro_var_;
            double ea = options.bias_acce_var_;

            double ev2 = ev;  // * ev;
            double et2 = et;  // * et;
            double eg2 = eg;  // * eg;
            double ea2 = ea;  // * ea;

            // 设置过程噪声
            Q_.diagonal() << 0, 0, 0, ev2, ev2, ev2, et2, et2, et2, eg2, eg2, eg2, ea2, ea2, ea2, 0, 0, 0;

            // 设置里程计噪声
            double o2 = options_.odom_var_ * options_.odom_var_;
            odom_noise_.diagonal() << o2, o2, o2;

            // 设置GNSS状态
            double gp2 = options.gnss_pos_noise_ * options.gnss_pos_noise_;
            double gh2 = options.gnss_height_noise_ * options.gnss_height_noise_;
            double ga2 = options.gnss_ang_noise_ * options.gnss_ang_noise_;
            gnss_noise_.diagonal() << gp2, gp2, gh2, ga2, ga2, ga2;
        }
        
        /// 更新名义状态变量，重置error state
        void UpdateAndReset() {
            p_ += dx_.template block<3, 1>(0, 0);
            v_ += dx_.template block<3, 1>(3, 0);
            R_ = R_ * SO3::exp(dx_.template block<3, 1>(6, 0));

            if (options_.update_bias_gyro_) {
                bg_ += dx_.template block<3, 1>(9, 0);
            }

            if (options_.update_bias_acce_) {
                ba_ += dx_.template block<3, 1>(12, 0);
            }

            g_ += dx_.template block<3, 1>(15, 0);

            ProjectCov();
            dx_.setZero();
        }

        /// 对P阵进行投影，参考式(3.63)
        void ProjectCov() {
            Mat18T J = Mat18T::Identity();
            J.template block<3, 3>(6, 6) = Mat3T::Identity() - 0.5 * SO3::hat(dx_.template block<3, 1>(6, 0));
            cov_ = J * cov_ * J.transpose();
        }
        
    private:
        /// 成员变量
        double current_time_ = 0.0;  // 当前时间

        /// 名义状态
        VecT p_ = VecT::Zero();
        VecT v_ = VecT::Zero();
        SO3 R_;
        VecT bg_ = VecT::Zero();
        VecT ba_ = VecT::Zero();
        VecT g_{0, 0, -9.8};

        /// 误差状态
        Vec18T dx_ = Vec18T::Zero();

        /// 协方差阵
        Mat18T cov_ = Mat18T::Identity();

        /// 噪声阵
        MotionNoiseT Q_ = MotionNoiseT::Zero();
        OdomNoiseT odom_noise_ = OdomNoiseT::Zero();
        GnssNoiseT gnss_noise_ = GnssNoiseT::Zero();
        LidarNoiseT lidar_noise = LidarNoiseT::Zero();
        /// 标志位
        bool first_gnss_ = true;  // 是否为第一个gnss数据
        bool first_odom_ = true;  // 是否为第一个odom数据
        /// 配置项
        EskfOptions options_;
};

using ESKFD = ESKF<double>;
using ESKFF = ESKF<float>;

template <typename S>
bool ESKF<S>::Predict(const IMU& imu) 
{
    assert(imu.timestamp_ >= current_time_);

    double dt = imu.timestamp_ - current_time_;
    if (dt > (5 * options_.imu_dt_) || dt < 0) {
        // 时间间隔不对，可能是第一个IMU数据，没有历史信息
        LOG(INFO) << "skip this imu because dt_ = " << dt;
        current_time_ = imu.timestamp_;
        return false;
    }

    // nominal state 递推
    // Vec3d world_imu = R_ * (imu.acce_ - ba_) + g_;
    // LOG(INFO) << world_imu;
    VecT new_p = p_ + v_ * dt + 0.5 * (R_ * (imu.acce_ - ba_)) * dt * dt + 0.5 * g_ * dt * dt;
    VecT new_v = v_ + R_ * (imu.acce_ - ba_) * dt + g_ * dt;
    SO3 new_R = R_ * SO3::exp((imu.gyro_ - bg_) * dt);

    R_ = new_R;
    v_ = new_v;
    p_ = new_p;
    // 其余状态维度不变   

    // error state 递推
    // 计算运动过程雅可比矩阵 F，见(3.47)
    // F实际上是稀疏矩阵，也可以不用矩阵形式进行相乘而是写成散装形式，这里为了教学方便，使用矩阵形式
    Mat18T F = Mat18T::Identity();                                                 // 主对角线
    F.template block<3, 3>(0, 3) = Mat3T::Identity() * dt;                         // p 对 v
    F.template block<3, 3>(3, 6) = -R_.matrix() * SO3::hat(imu.acce_ - ba_) * dt;  // v对theta
    F.template block<3, 3>(3, 12) = -R_.matrix() * dt;                             // v 对 ba
    F.template block<3, 3>(3, 15) = Mat3T::Identity() * dt;                        // v 对 g
    F.template block<3, 3>(6, 6) = SO3::exp(-(imu.gyro_ - bg_) * dt).matrix();     // theta 对 theta
    F.template block<3, 3>(6, 9) = -Mat3T::Identity() * dt;

        // mean and cov prediction
    dx_ = F * dx_;  // 这行其实没必要算，dx_在重置之后应该为零，因此这步可以跳过，但F需要参与Cov部分计算，所以保留
    cov_ = F * cov_.eval() * F.transpose() + Q_;
    current_time_ = imu.timestamp_;
    
    return true;
}

template <typename S>
bool ESKF<S>::ObserveLidar(const PoseData& pos)
{
    SE3 se3; pos.GetSE3(se3);
    ObserveSE3(se3);
}

template <typename S>
bool ESKF<S>::ObserveSE3(const SE3& pose, double trans_noise, double ang_noise) 
{
    /// 既有旋转，也有平移
    /// 观测状态变量中的p, R，H为6x18，其余为零
    Eigen::Matrix<S, 6, 18> H = Eigen::Matrix<S, 6, 18>::Zero();
    H.template block<3, 3>(0, 0) = Mat3T::Identity();  // P部分
    H.template block<3, 3>(3, 6) = Mat3T::Identity();  // R部分（3.66)

    // 卡尔曼增益和更新过程
    Vec6d noise_vec;
    noise_vec << trans_noise, trans_noise, trans_noise, ang_noise, ang_noise, ang_noise;

    Mat6d V = noise_vec.asDiagonal();
    Eigen::Matrix<S, 18, 6> K = cov_ * H.transpose() * (H * cov_ * H.transpose() + V).inverse();

    // 更新x和cov
    Vec6d innov = Vec6d::Zero();
    innov.template head<3>() = (pose.translation() - p_);          // 平移部分
    innov.template tail<3>() = (R_.inverse() * pose.so3()).log();  // 旋转部分(3.67)

    dx_ = K * innov;
    cov_ = (Mat18T::Identity() - K * H) * cov_;

    UpdateAndReset();
    return true;
}
}
